import os
import numpy as np
import pylab as plt
from params import *

datasets = ['bkite','fsquare','gowalla']
lbls = ['Brightkite', 'Foursquare', 'Gowalla']
node_dist = [5651.0,8494.0,5663.0]
link_dist = [2041.0,1442.0,1792.0]
degs = [7.88,22.07,9.48,7.0]

def compute_distance_degree(graph_file):
    data = {}
    def update_node(n,d):
        if not n in data:
            data[n] = (0,0.0)
        deg,tot_dist = data[n]
        deg += 1
        tot_dist += d
        data[n] = (deg,tot_dist)

    avg_dist = 0.0
    k = 0
    for line in open(graph_file):
        k += 1
        n1,n2,d = line.strip().split()
        n1,n2 = map(int,(n1,n2))
        d =float(d)
        avg_dist += d
        update_node(n1,d)
        update_node(n2,d)

    avg_dist = avg_dist / k
    print "read %d nodes"%(len(data))
    print 'avg link length ', avg_dist

    values = []
    for n in data:
        deg,tot_dist = data[n]
        values.append(tot_dist/deg)

    return values

plots = []
for dataset,m in zip(datasets,symbols):
    file = os.path.join(WORKDIR, 'gsn','results',dataset,'%s_distance_strength.txt'%dataset)
    file = os.path.join(WORKDIR, 'gsn','traces',dataset,'%s_graph.txt'%dataset)
    filegeo = os.path.join(WORKDIR, 'gsn','null_graphs',dataset,'%s_geo_null_graph_2.txt'%dataset)
    filesocial = os.path.join(WORKDIR, 'gsn','null_graphs',dataset,'%s_social_null_graph_2.txt'%dataset)
    print file
#    data = []
#    k = 0
#    c = 0
#    for line in open(file):
#        deg,strength = map(float,line.strip().split(';'))
#        if strength/deg < 100:
#            c += 1
#        m = max(0.0, strength/deg)
#        data.append(m)
#        k += 1
#    print k,c, 1.0*c/k
    data = compute_distance_degree(file)
    datageo = compute_distance_degree(filegeo)
    datasocial  = compute_distance_degree(filesocial)

    def plot_cdf(data):
        bb = np.logspace(0,5)
        hp,hb,xx = plt.hist(data,
                bins=bb)
                #cumulative=F,normed=True)#,log=True)
        hc = [hp[0]]
        for v in hp[1:]:
            hc.append(v+hc[-1])
        tot = hc[-1]
        hc = [1.0*i/tot for i in hc]
        return hb[:-1],hc
    x,y = plot_cdf(data)
    xg,yg = plot_cdf(datageo)
    xs,ys = plot_cdf(datasocial)

    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    plt.semilogx(x,y,'k-')
    plt.semilogx(xg,yg,'k:')
    plt.semilogx(xs,ys,'k--')
    plt.legend(['Original data', 'Geo model', 'Social model'],
            loc='upper left')
    plt.grid(True)
    plt.xlabel('Average friend distance [km]')
    plt.xlabel(r'$w$ [km]')
    plt.ylabel('CDF')
    #plt.xlim(xmax=1e4)
    plt.ylim(ymax=1)
    plt.savefig('%s_dist_strength_cdf.pdf'%dataset)
    plt.close()
